load(
    "//tensorflow/lite/micro:build_def.bzl",
    "tflm_cc_library",
    "tflm_cc_test",
)
load(
    "@flatbuffers//:build_defs.bzl",
    "flatbuffer_cc_library",
    "flatbuffer_py_library",
)
load("@rules_python//python:defs.bzl", "py_test")
load("@tflm_pip_deps//:requirements.bzl", "requirement")

package(
    default_visibility = [
        "//visibility:public",
    ],
)

flatbuffer_cc_library(
    # Generates the header-only library "metadata_generated.h", used to read
    # the metadata flatbuffer.
    name = "metadata_cc",
    srcs = ["metadata.fbs"],
)

tflm_cc_library(
    # The header-only library generated by flatc in ":metadata_cc" is saved to
    # the source tree and comitted to git as "metadata_saved.h", which is used
    # by code which builds via the Make build system, which has no means of
    # generating the header on the fly. Code which builds via both bazel and
    # Make should #include the saved header and use this target in its bazel
    # BUILD deps. Code built exclusively via bazel would typically depend
    # directly on ":metadata_cc", which would generate a header from the schema
    # on the fly, during the build.
    #
    # When the schema definition "metadata.fbs" is changed, this saved header
    # should be updated by running the script "./metadata_saved_update.sh",
    # outside of bazel (because bazel cannot modify the source tree). The
    # script regenerates the header from the schema and copies it to the source
    # tree as "metadata_saved.h".
    #
    # Comitting the generated file risks inconsistency between the schema and
    # the saved header, so consistency ensured by the unit test
    # ":metadata_saved_test".
    #
    name = "metadata_saved",
    hdrs = ["metadata_saved.h"],
)

sh_test(
    # Ensures consistency bewteen the schema and the saved generated header.
    # Fails if they mismatch, in which case, ./metadata_saved_update.sh should
    # be run. See :metadata_saved above.
    name = "metadata_saved_test",
    size = "small",
    srcs = ["metadata_saved_test.sh"],
    args = [
        "$(location metadata_saved.h)",
        "$(location :metadata_cc_srcs)",
    ],
    data = [
        "metadata_saved.h",
        ":metadata_cc_srcs",
    ],
)

tflm_cc_test(
    name = "metadata_test_cc",
    size = "small",
    srcs = ["metadata_test.cc"],
    deps = [
        ":metadata_saved",
        "//tensorflow/lite/micro:hexdump",
        "//tensorflow/lite/micro/testing:micro_test",
        "@flatbuffers//:runtime_cc",
    ],
)

flatbuffer_py_library(
    # Generates the Python module "metadata_py_generated", used to read the
    # metadata flatbuffer.
    name = "metadata_py",
    srcs = ["metadata.fbs"],
)

py_test(
    name = "metadata_test_py",
    size = "small",
    srcs = ["metadata_test.py"],
    main = "metadata_test.py",
    deps = [
        "metadata_py",
        "@flatbuffers//:runtime_py",
        requirement("tensorflow"),
    ],
)

py_binary(
    name = "compress",
    srcs = [
        "compress.py",
    ],
    deps = [
        ":metadata_py",
        ":model_facade",
        ":spec",
        "@absl_py//absl:app",
        "@absl_py//absl/flags",
        "@flatbuffers//:runtime_py",
        requirement("bitarray"),
        requirement("numpy"),
    ],
)

py_test(
    name = "compress_test",
    size = "small",
    srcs = [
        "compress_test.py",
    ],
    deps = [
        ":compress",
        ":metadata_py",
        ":model_facade",
        ":spec",
        ":test_models",
        "//tensorflow/lite/python:schema_py",
        requirement("bitarray"),
        requirement("numpy"),
        requirement("tensorflow"),
    ],
)

py_library(
    name = "model_facade",
    srcs = ["model_facade.py"],
    deps = [
        "//tensorflow/lite/python:schema_py",
        requirement("flatbuffers"),
    ],
)

py_test(
    name = "model_facade_test",
    size = "small",
    srcs = ["model_facade_test.py"],
    deps = [
        ":model_facade",
        ":test_models",
        requirement("tensorflow"),
    ],
)

py_library(
    name = "spec",
    srcs = ["spec.py"],
    deps = [
        requirement("pyyaml"),
    ],
)

py_test(
    name = "spec_test",
    size = "small",
    srcs = ["spec_test.py"],
    deps = [
        ":spec",
        requirement("tensorflow"),
    ],
)

py_library(
    name = "test_models",
    srcs = ["test_models.py"],
    deps = [
        "//tensorflow/lite/python:schema_py",
        requirement("flatbuffers"),
        requirement("numpy"),
    ],
)

py_test(
    name = "test_models_test",
    size = "small",
    srcs = ["test_models_test.py"],
    deps = [
        ":test_models",
        "//tensorflow/lite/python:schema_py",
        requirement("tensorflow"),
    ],
)
